package memoryx

import (
	"context"
	"sync"

	"github.com/tmc/langchaingo/schema"

	"gitee.com/dn-jinmin/langchaingo"
)

type Memoryx struct {
	sync.Mutex
	getMemory     func() schema.Memory
	memorys       map[string]schema.Memory
	defaultMemory schema.Memory
}

func NewMemoryx(handle func() schema.Memory) *Memoryx {
	return &Memoryx{
		defaultMemory: handle(),
		getMemory:     handle,
		memorys:       make(map[string]schema.Memory),
	}
}

func (m *Memoryx) GetMemoryKey(ctx context.Context) string {
	return m.defaultMemory.GetMemoryKey(ctx)
}

func (m *Memoryx) MemoryVariables(ctx context.Context) []string {
	return m.defaultMemory.MemoryVariables(ctx)
}

func (m *Memoryx) LoadMemoryVariables(ctx context.Context, inputs map[string]any) (map[string]any, error) {
	return m.memory(ctx).LoadMemoryVariables(ctx, inputs)
}

func (m *Memoryx) SaveContext(ctx context.Context, inputs map[string]any, outputs map[string]any) error {
	return m.memory(ctx).SaveContext(ctx, inputs, outputs)
}

func (m *Memoryx) Clear(ctx context.Context) error {
	return m.memory(ctx).Clear(ctx)
}

func (m *Memoryx) memory(ctx context.Context) schema.Memory {
	m.Lock()
	defer m.Unlock()

	var chatId string
	v := ctx.Value(langchaingo.ChatId)

	if v == nil {
		return m.defaultMemory
	}

	chatId = v.(string)
	cm, ok := m.memorys[chatId]
	if !ok {
		cm = m.getMemory()
		m.memorys[chatId] = cm
	}
	return cm
}
